package com.tlzh.controller;



import com.fasterxml.jackson.databind.ObjectMapper;
import com.tlzh.pojo.dto.ChatRequest;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;


import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@RequestMapping("/chatV2")
@CrossOrigin
public class ChatV2Controller {

    @Resource
    private  WebClient webClient;
    // 存储每个用户的中断状态，使用 ConcurrentHashMap 保证线程安全
    private final Map<String, Boolean> interruptionStatusMap = new ConcurrentHashMap<>();
    // 存储每个用户的 SseEmitter，使用 ConcurrentHashMap 保证线程安全
    private final Map<String, SseEmitter> emitterMap = new ConcurrentHashMap<>();


    @PostMapping(value = "/ai/ollamaChatStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatModel(@RequestBody ChatRequest request ) {

        String token = request.getToken();

        SseEmitter emitter = new SseEmitter();
        interruptionStatusMap.put(token, false);
        emitterMap.put(token, emitter);

        // 实际的 API 地址，你需要根据实际情况修改
        String apiUrl = "http://192.168.0.188:11434/api/chat";

        // logger 获取 请求时间
        System.out.println("Request received at: " + LocalDateTime.now() + " for token: " + token);
        // 使用 WebClient 发送异步 POST 请求并获取 Flux<String>
        Flux<String> responseFlux = webClient.post()
                .uri(apiUrl)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class);
//        System.out.println("Request Endt: " + LocalDateTime.now() + " for token: " + token);

//        System.out.println("Mono<Void> mono: " + LocalDateTime.now() + " for token: " + token);
        // 订阅 Flux 并将数据发送到 SseEmitter
        Mono<Void> mono = responseFlux
                .takeWhile(part -> !interruptionStatusMap.getOrDefault(token, false)) // 当该用户的中断状态为 true 时停止发送数据
                .doOnNext(part -> {
                    try {

//                        ObjectMapper objectMapper = new ObjectMapper();
//                        JsonNode jsonNode = objectMapper.readTree(part);
//
//                        // 提取 context 数据
//                        JsonNode contextNode = jsonNode.get("message").get("content");
//                        String context = contextNode.asText();
//                        System.out.println(LocalDateTime.now()+ "context:{"+ context + "} for token: " + token);
                        SseEmitter.SseEventBuilder event = SseEmitter.event()
                                .data(part)
                                .id(String.valueOf(System.currentTimeMillis()))
                                .reconnectTime(50000);
//                        System.out.println("event Start: " + LocalDateTime.now() + " for token: " + token);
                        emitter.send(event);
//                        System.out.println("event End: " + LocalDateTime.now() + " for token: " + token);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                })
                .doOnComplete(() -> {
                    emitter.complete();
                    interruptionStatusMap.remove(token);
                    emitterMap.remove(token);
                })
                .doOnError(error -> {
                    emitter.completeWithError(error);
                    interruptionStatusMap.remove(token);
                    emitterMap.remove(token);
                })
                .then();

        // 启动异步处理
        mono.subscribe();

        return emitter;
    }

    @PostMapping("/ai/interruptOllamaChatStream")
    public void interruptStreaming(@RequestHeader("token") String token) {
        interruptionStatusMap.put(token, true);
        SseEmitter emitter = emitterMap.get(token);
        if (emitter != null) {
            emitter.complete();
            emitterMap.remove(token);
            interruptionStatusMap.remove(token);
        }
    }


    private final ObjectMapper objectMapper = new ObjectMapper();
    private static final Logger logger = LoggerFactory.getLogger(FluxController.class);

    private final Map<String, Disposable> activeRequests = new ConcurrentHashMap<>();

    // 存储每个用户的中断状态
    private  Map<String, AtomicBoolean> interruptionStatusMapV2 = new ConcurrentHashMap<>();


    @PostMapping(value = "/ai/ollamaStreamV1", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> queryModelV1(@RequestBody ChatRequest request) {
        String uuid = request.getToken();
        try {
            String requestBody = objectMapper.writeValueAsString(request);
            logger.info("Sending request body: {}", requestBody);
        } catch (Exception e) {
            logger.error("Error converting request to JSON", e);
        }

        // 初始化用户的中断状态
        AtomicBoolean isInterrupted = interruptionStatusMapV2.computeIfAbsent(uuid, k -> new AtomicBoolean(false));


        Flux<ServerSentEvent<String>> flux = webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .header("Accept-Charset", "UTF-8")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class)
                .map(response -> ServerSentEvent.<String>builder().data(response).build())
                .takeWhile(part -> !isInterrupted.get())
                .doFinally(signal -> {
                    interruptionStatusMapV2.remove(uuid);
                    activeRequests.remove(uuid);
                });


        // 订阅 Flux 并保存 Disposable 到 activeRequests 中
        Disposable disposable = flux.subscribeOn(Schedulers.boundedElastic())
                .doFinally(signalType -> activeRequests.remove(uuid))
                .subscribe();

        activeRequests.put(uuid, disposable);

        return flux;
    }

    // 提供一个接口用于根据 UUID 打断请求
    @GetMapping("/ai/interrupt")
    public void interruptRequest(String uuid) {
        Disposable disposable = activeRequests.get(uuid);
        AtomicBoolean isInterrupted = interruptionStatusMapV2.get(uuid);
        if (disposable != null) {
            isInterrupted.set(true);
            logger.info("Request with UUID {} has been interrupted.", uuid);
        }
        // 可选：移除中断状态和Disposable记录
        interruptionStatusMapV2.remove(uuid);
        activeRequests.remove(uuid);
    }



    /**
    * @新增一个 客户端 断开 SSE 服务端 Flux 也断开
    * */
    @PostMapping(value = "/ai/ollamaStreamV2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> queryModelV2(@RequestBody ChatRequest request, ServerWebExchange exchange) {
        String uuid = request.getToken();
        try {
            String requestBody = objectMapper.writeValueAsString(request);
            logger.info("Sending request body: {}", requestBody);
        } catch (Exception e) {
            logger.error("Error converting request to JSON", e);
        }

        // 初始化用户的中断状态
        AtomicBoolean isInterrupted = interruptionStatusMapV2.computeIfAbsent(uuid, k -> new AtomicBoolean(false));

        // 获取 ServerHttpResponse
        ServerHttpResponse response = exchange.getResponse();

        // 监听响应关闭事件
//        Mono<Void> onResponseClose = response.beforeCommit(() -> {
//            if (response.isCommitted()) {
//                // 设置中断状态
//                isInterrupted.set(true);
//            }
//            return Mono.empty();
//        });

        Flux<ServerSentEvent<String>> flux = webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .header("Accept-Charset", "UTF-8")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class)
                .map(responseData -> ServerSentEvent.<String>builder().data(responseData).build())
                .takeWhile(part -> !isInterrupted.get())
                .doFinally(signal -> interruptionStatusMapV2.remove(uuid));

        // 合并 Flux 和响应关闭事件
//        return Flux.merge(flux, onResponseClose.thenMany(Flux.empty()));

        return flux;
    }

}
